package m3f.io;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Random;

import m3f.matrix.Minibatch;
import m3f.matrix.SparseVector;

public class CachedMinibatchLoader implements Runnable {
	
	private MatrixReader mainReader;
	private String filename;
	private int batchPointer;
	private int totalReadRows;
	private ArrayList<SparseVector> batch;
	private ArrayList<SparseVector> cache;
	private boolean reading;
	private int vectorSize;
	private int matrixRows;
	private int globalPointer;
	private int cacheSize;
	private long randomSeed;
	private boolean normalizeVectors;
	
	public CachedMinibatchLoader(String file, boolean normalize){
		this.filename = file;
		this.normalizeVectors = normalize;
		cache = null;
		batchPointer = 0;
		totalReadRows = 0;
		mainReader = null;
		reading = true;
		vectorSize = 0;
		matrixRows = 0;
		globalPointer = 0;
		cacheSize = 10000;
	}
	
	private synchronized void startReading(){
		reading = true;
		notifyAll();
	}
	
	private synchronized void stopReading(){
		reading = false;
		notifyAll();
	}
	
	public int getCacheSize(){
		return cacheSize;
	}
	
	public int getMatrixRows(){
		if(matrixRows == 0){
			waitWhileReading();
		}
		return matrixRows;
	}
	
	public int getVectorSize(){
		if(vectorSize == 0){
			waitWhileReading();
		}
		return vectorSize;
	}
	
	public boolean hasData(){
		if(matrixRows == 0){
			waitWhileReading();
		}
		return totalReadRows < matrixRows;
	}
	
	private ArrayList<SparseVector> loadBatch(){
		ArrayList<SparseVector> read = mainReader.readVectors( cacheSize );
		if(randomSeed != 0){
			Collections.shuffle(read, new Random(randomSeed));
		}
		globalPointer += read.size();
		return read;
	}
	
	private void loadCacheInBackground(){
		startReading();
		if(cache == null && batchPointer >= batch.size()/2 && globalPointer < matrixRows){
			(new Thread(){
				@Override
				public void run(){
					cache = loadBatch();
					stopReading();
				}
			}).start();
		}else{
			stopReading();
		}
	}
	
	public synchronized Minibatch nextMinibatch(int minibatch){
		waitWhileReading();
		int mbSize = minibatch;
		if(batchPointer + minibatch > batch.size() && totalReadRows + minibatch > matrixRows){
			mbSize = batch.size() - batchPointer;
		}
		int id = (int)Math.floor(((double)(totalReadRows)/(double)minibatch));
		if (totalReadRows == matrixRows) {
			return null;
		}
		Minibatch mb = new Minibatch(vectorSize, mbSize, id);
		int insertedCols;
		for(insertedCols = 0; insertedCols < mbSize && batchPointer < batch.size(); insertedCols++){
			mb.setColumn(insertedCols, batch.get(batchPointer));
			batchPointer++;
			totalReadRows++;
		}
		// Check cache state
		loadCacheInBackground();
		if(batchPointer == batch.size()){
			if(cache != null){
				batch = cache;
				cache = null;
			}
			batchPointer = 0;
		}
		// Load pending vectors for this minibatch
		int pendingVectors = mbSize - insertedCols;
		if(pendingVectors > 0){
			waitWhileReading();
			for(int i = 0; i < pendingVectors && batchPointer < batch.size(); i++){
				mb.setColumn(insertedCols+i, batch.get(batchPointer));
				batchPointer++;
				totalReadRows++;
			}
		}
		return mb;
	}
	
	public String toString(){
		String s = "CachedMinibatchReader: (" + filename + "=> read:" + totalReadRows + "; pointer:" + globalPointer + "; total:" + matrixRows + ")";
		if(batch == null) s += "[MainData=>null] ";
		else{ 
			s += "[MainData=> total:" + batch.size();
			s += "; pointer: "+ batchPointer + "]";
		}
		if(cache == null) s+= "[Cache=>null] ";
		else s += "[Cache=>"+cache.size()+"]";
		return s;
	}
	
	public synchronized void reset(long randomSeed){
		waitWhileReading();
		startReading();
		this.randomSeed = randomSeed;
		(new Thread(){
			@Override
			public void run(){
				batch = loadBatch();
				stopReading();
			}
		}).start();
		globalPointer = 0;
		batchPointer = 0;
		totalReadRows = 0;
		cache = null;
	}
	
	@Override
	public void run(){
		startReading();
		mainReader = MatrixReaderFactory.newInstance(filename, normalizeVectors);
		System.out.print(filename);
		mainReader.start(filename);
		vectorSize = mainReader.getMatrixColumns();
		matrixRows = mainReader.getMatrixRows();
		stopReading();
	}
	
	public void setCacheSize(int cache){
		cacheSize = cache;
	}
	
	private synchronized void waitWhileReading(){
		while(reading){
			try{
				wait();
			}catch(InterruptedException e){}
		}		
	}
}
